#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
链路追踪服务
"""

import logging
from datetime import datetime
from bson import ObjectId
from db.mongo import get_db
from base.models.trace import TraceModel, ServiceModel, TopologyModel, RequestModel

logger = logging.getLogger(__name__)

class TraceService:
    """链路追踪服务"""

    @staticmethod
    def create_trace(request_definition, trace_data):
        """
        创建链路追踪记录

        Args:
            request_definition: 请求定义
            trace_data: 链路数据

        Returns:
            str: 链路ID
        """
        db = get_db()

        # 创建链路记录
        trace = TraceModel(request_definition=request_definition, trace_data=trace_data)
        trace_dict = trace.to_dict()

        # 插入链路记录
        db.traces.insert_one(trace_dict)
        trace_id = trace_dict["_id"]

        # 处理服务和拓扑关系
        service_map = {}  # 用于存储服务名称到服务ID的映射

        for span in trace_data:
            # 尝试从不同的字段中提取服务名称
            service_name = None
            if 'service' in span:
                if isinstance(span['service'], dict) and 'name' in span['service']:
                    service_name = span['service']['name']
                elif isinstance(span['service'], str):
                    service_name = span['service']
            elif 'service_name' in span:
                service_name = span['service_name']
            elif 'service.name' in span.get('tags', {}):
                service_name = span['tags']['service.name']
            elif 'attributes' in span and 'service.name' in span['attributes']:
                service_name = span['attributes']['service.name']

            span_id = span.get("id")
            parent_id = span.get("parent_id")

            # 如果服务名称为空，则跳过
            if not service_name:
                continue

            # 如果服务不存在，则创建服务记录
            if service_name not in service_map:
                service = ServiceModel(trace_id=trace_id, service_name=service_name)
                service_dict = service.to_dict()
                db.trace_services.insert_one(service_dict)
                service_id = service_dict["_id"]
                service_map[service_name] = service_id
            else:
                service_id = service_map[service_name]

            # 创建拓扑关系
            if parent_id:
                # 查找父span所属的服务
                parent_service = None
                for parent_span in trace_data:
                    if parent_span.get("id") == parent_id:
                        # 尝试从不同的字段中提取父服务名称
                        if 'service' in parent_span:
                            if isinstance(parent_span['service'], dict) and 'name' in parent_span['service']:
                                parent_service = parent_span['service']['name']
                            elif isinstance(parent_span['service'], str):
                                parent_service = parent_span['service']
                        elif 'service_name' in parent_span:
                            parent_service = parent_span['service_name']
                        elif 'service.name' in parent_span.get('tags', {}):
                            parent_service = parent_span['tags']['service.name']
                        elif 'attributes' in parent_span and 'service.name' in parent_span['attributes']:
                            parent_service = parent_span['attributes']['service.name']
                        break

                if parent_service and parent_service in service_map:
                    parent_service_id = service_map[parent_service]
                    # 创建拓扑关系记录
                    topology = TopologyModel(service_id=service_id, parent_id=parent_service_id)
                    db.trace_topology.insert_one(topology.to_dict())

            # 处理请求信息
            attributes = span.get("attributes", {})

            # 处理入口请求 (ingress)
            if not parent_id or parent_id == "":
                method = attributes.get("http.request.method")
                path = attributes.get("url.path") or attributes.get("http.route")

                if method and path:
                    request = RequestModel(
                        service_id=service_id,
                        request_type="ingress",
                        method=method,
                        url=path
                    )
                    db.trace_requests.insert_one(request.to_dict())

            # 处理出口请求 (egress)
            if "url.full" in attributes:
                method = attributes.get("http.request.method")
                url = attributes.get("url.full")

                if method and url:
                    request = RequestModel(
                        service_id=service_id,
                        request_type="egress",
                        method=method,
                        url=url
                    )
                    db.trace_requests.insert_one(request.to_dict())

        return trace_id

    @staticmethod
    def get_traces(page=1, page_size=10):
        """
        获取链路列表

        Args:
            page: 页码
            page_size: 每页大小

        Returns:
            dict: 包含链路列表和总数的字典
        """
        db = get_db()

        # 计算跳过的记录数
        skip = (page - 1) * page_size

        # 查询链路记录
        traces = list(db.traces.find().sort("created_at", -1).skip(skip).limit(page_size))

        # 获取总记录数
        total = db.traces.count_documents({})

        # 计算每个链路的服务数量和span数量
        for trace in traces:
            trace_id = trace["_id"]

            # 计算服务数量
            service_count = db.trace_services.count_documents({"trace_id": trace_id})
            trace["service_count"] = service_count

            # 计算span数量
            trace["span_count"] = len(trace.get("trace_data", []))

            # 移除大型数据字段以减少响应大小
            if "trace_data" in trace:
                del trace["trace_data"]

        return {
            "traces": traces,
            "total": total,
            "page": page,
            "page_size": page_size,
            "pages": (total + page_size - 1) // page_size
        }

    @staticmethod
    def get_trace_by_id(trace_id):
        """
        根据ID获取链路详情

        Args:
            trace_id: 链路ID

        Returns:
            dict: 链路详情
        """
        db = get_db()

        # 查询链路记录
        trace = db.traces.find_one({"_id": trace_id})
        if not trace:
            return None

        return trace

    @staticmethod
    def get_trace_services(trace_id):
        """
        获取链路相关的服务列表

        Args:
            trace_id: 链路ID

        Returns:
            list: 服务列表
        """
        db = get_db()

        # 查询服务记录
        services = list(db.trace_services.find({"trace_id": trace_id}))

        return services

    @staticmethod
    def get_trace_topology(trace_id):
        """
        获取链路拓扑关系

        Args:
            trace_id: 链路ID

        Returns:
            list: 拓扑关系列表
        """
        db = get_db()

        # 查询服务记录
        services = list(db.trace_services.find({"trace_id": trace_id}))

        # 创建服务ID到服务名称的映射
        service_map = {service["_id"]: service for service in services}

        # 查询拓扑关系
        topology_list = []

        for service in services:
            service_id = service["_id"]

            # 查询当前服务的拓扑关系
            topologies = list(db.trace_topology.find({"service_id": service_id}))

            for topology in topologies:
                parent_id = topology["parent_id"]
                if parent_id in service_map:
                    topology_list.append({
                        "source": service_map[parent_id]["service_name"],
                        "target": service["service_name"]
                    })

        return topology_list

    @staticmethod
    def get_service_requests(service_id):
        """
        获取服务的请求信息

        Args:
            service_id: 服务ID

        Returns:
            dict: 包含入口和出口请求的字典
        """
        db = get_db()

        # 查询入口请求
        ingress_requests = list(db.trace_requests.find({
            "service_id": service_id,
            "request_type": "ingress"
        }))

        # 查询出口请求
        egress_requests = list(db.trace_requests.find({
            "service_id": service_id,
            "request_type": "egress"
        }))

        return {
            "ingress": ingress_requests,
            "egress": egress_requests
        }

    @staticmethod
    def delete_trace(trace_id):
        """
        删除链路记录

        Args:
            trace_id: 链路ID

        Returns:
            bool: 是否成功删除
        """
        db = get_db()

        # 查询链路记录
        trace = db.traces.find_one({"_id": trace_id})
        if not trace:
            return False

        # 查询相关服务
        services = list(db.trace_services.find({"trace_id": trace_id}))
        service_ids = [service["_id"] for service in services]

        # 删除拓扑关系
        db.trace_topology.delete_many({"service_id": {"$in": service_ids}})
        db.trace_topology.delete_many({"parent_id": {"$in": service_ids}})

        # 删除请求信息
        db.trace_requests.delete_many({"service_id": {"$in": service_ids}})

        # 删除服务记录
        db.trace_services.delete_many({"trace_id": trace_id})

        # 删除链路记录
        db.traces.delete_one({"_id": trace_id})

        return True
